#pragma once
#include "base_detector.h"
#include "TopsInference/TopsInferRuntime.h"
#include <numeric>


const char* DataTypeString(int dtype);

struct ShapeInfo
{
    std::string name;
    std::vector<int> dims;
    int dtype;
    int dtype_size;
    int volume;
    int mem_size;
    ShapeInfo() {}
    ShapeInfo(const char *tensor_name, std::vector<int> &_dims, int dtype, int _dtype_size, int _mem_size)
    : name(tensor_name), dims(_dims), dtype(dtype), dtype_size(_dtype_size), mem_size(_mem_size) {
        volume = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<int>());
    }

    std::string toString()
    {
        std::stringstream ss;
        ss << name << ": [";
        for (size_t i = 0; i < dims.size(); ++i)
        {
            ss << dims[i];
            if (i < dims.size() - 1)
                ss << ",";
        }
        ss << "] ";
        ss << DataTypeString(dtype) << ",";
        ss << "mem_size:" << mem_size;
        return ss.str();
    }
};

class Yolo11GcuDetector : public BaseDetector {
public:
    Yolo11GcuDetector(const std::string& model_path, int dev_id, int clusterid);
    ~Yolo11GcuDetector() override;
    std::vector<DetectionResult> detect(const cv::Mat& image) override;
    bool isInitialized() const override { return m_initialized; }
protected:
    int loadModel(const std::string& model_file);
    std::vector<ShapeInfo> getInputsShape();
    std::vector<ShapeInfo> getOutputsShape();
    int preProcess(std::vector<cv::Mat> &images);
    int postProcess(std::vector<cv::Mat> &images, std::vector<std::vector<DetectionResult>> &batchDets);

protected:
    bool m_initialized {false};
    float m_confThres = 0.25;
    float m_nmsThres = 0.45;
    TopsInference::handler_t mTopsHandle = nullptr;    
    TopsInference::IEngine *mTopsEngine;    
    std::vector<void*> mNetInputs;    
    std::vector<void*> mNetOutputs;    
    std::vector<ShapeInfo> mInputShapes;    
    std::vector<ShapeInfo> mOutputShapes;
    int m_devId;
    int m_clusterId;
    std::string m_model_path;
};
